package com.qjc.demo.train.redblacktree;

/**
 * @ClassName: RBTree
 * @Description: 红黑树
 * @Author: qjc
 * @Date: 2022/5/1 7:11 PM
 */
public class RBTree {
    RBTreeNode root;

    // 遍历
    public static void list(RBTreeNode node) {
        if (node == null) return;
        // 递归终止条件，叶子结点
        if (node.getLeft() == null && node.getRight() == null) {
            System.out.println(node);
            return;
        }
        System.out.println(node);
        list(node.getLeft());
        list(node.getRight());
    }

    public void insert(int key) {
        RBTreeNode node = new RBTreeNode(key);
        if (root == null) {
            // 根节点是黑色的
            node.setBlack(true);
            root = node;
            return;
        }
        RBTreeNode parent = root;
        RBTreeNode son = null;
        // 让son指向左节点/右节点
        if (key <= parent.getKey()) {
            son = parent.getLeft();
        } else {
            son = parent.getRight();
        }
        // 如果左节点/右节点不为空，则循环查找，直到找到左节点/右节点为空的节点
        while (son != null) {
            parent = son;// 父指针指向子节点
            if (key <= parent.getKey()) {
                son = parent.getLeft();
            } else {
                son = parent.getRight();
            }
        }
        // 指针知道了最后，开始赋值
        if (key <= parent.getKey()) {
            parent.setLeft(node);
        } else {
            parent.setRight(node);
        }
        node.setParent(parent);

        // 自平衡
        balanceInsert(node);
    }

    private void balanceInsert(RBTreeNode node) {
        RBTreeNode father, grandFather;
        // 父节点为红的
        while ((father = node.getParent()) != null && father.isBlack() == false) {
            // 父节点不为空，且是红色的，则进行自平衡
            grandFather = father.getParent();
            if (father == grandFather.getLeft()) {
                // 父节点在祖父节点的左边
                // 叔节点
                RBTreeNode uncle = grandFather.getRight();
                if (uncle != null && uncle.isBlack() == false) {
                    // 进行颜色反转
                    setBlack(father);
                    setBlack(uncle);
                    setRed(grandFather);
                    node = grandFather;
                    continue;
                }
                if (node == father.getRight()) {
                    // 进行左旋
                    leftRotate(father);

                    RBTreeNode tmp = node;
                    node = father;
                    father = tmp;
                }
                setBlack(father);
                setRed(grandFather);
                // 进行右旋
                rightRotate(grandFather);
            } else {
                // 父节点在祖父节点的右边
                // 叔节点
                RBTreeNode uncle = grandFather.getLeft();
                if (uncle != null && uncle.isBlack() == false) {
                    // 进行颜色反转
                    setBlack(father);
                    setBlack(uncle);
                    setRed(grandFather);
                    node = grandFather;
                    continue;
                }
                if (node == father.getLeft()) {
                    // 进行右旋
                    rightRotate(father);

                    RBTreeNode tmp = node;
                    node = father;
                    father = tmp;
                }
                setBlack(father);
                setRed(grandFather);
                // 进行左旋
                leftRotate(grandFather);
            }
        }
        setBlack(root);
    }

    // 左旋
    private void leftRotate(RBTreeNode node) {
        RBTreeNode right = node.getRight();
        RBTreeNode parent = node.getParent();
        if (parent == null) {
            root = right;
            right.setParent(null);
        } else {
            if (parent.getLeft() != null && parent.getLeft() == node) {
                parent.setLeft(right);
            } else {
                parent.setRight(right);
            }
            right.setParent(parent);
        }
        node.setParent(right);
        node.setRight(right.getLeft());
        if (right.getLeft() != null) {
            right.getLeft().setParent(node);
        }
        right.setLeft(node);
    }

    // 右旋
    private void rightRotate(RBTreeNode node) {
        RBTreeNode left = node.getLeft();
        RBTreeNode parent = node.getParent();
        if (parent == null) {
            root = left;
            left.setParent(null);
        } else {
            if (parent.getLeft() != null && parent.getLeft() == node) {
                parent.setLeft(left);
            } else {
                parent.setRight(left);
            }
            left.setParent(parent);
        }
        node.setParent(left);
        node.setLeft(left.getRight());
        if (left.getRight() != null) {
            left.getRight().setParent(node);
        }
        left.setRight(node);
    }

    private void setBlack(RBTreeNode node) {
        node.setBlack(true);
    }

    private void setRed(RBTreeNode node) {
        node.setBlack(false);
    }
    public static void main(String[] args) {
        RBTree rb=new RBTree();
        rb.insert(10);//根节点
        rb.insert(5);
        rb.insert(9);
        rb.insert(3);
        rb.insert(6);
        rb.insert(7);
        rb.insert(19);
        rb.insert(32);
        rb.insert(24);
        rb.insert(17);

        rb.list(rb.root);
    }

}
